import os

from langchain_community.vectorstores import FAISS, AzureSearch

from .base import VectorBase


#
class AzureSearchDatabase(VectorBase):

    def save_chunks_into_vectorstore(self, content_chunks, embedding_model):
        index_name: str = "langchain-vector-demo"
        vector_store = AzureSearch(
            azure_search_endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
            azure_search_key=os.environ["AZURE_SEARCH_KEY"],
            index_name=index_name,
            embedding_function=embedding_model.embed_query,
        )
        vector_store.add_documents(documents=content_chunks)
        return vector_store

    def get_vectorstore(self, embedding_model):
        index_name: str = "langchain-vector-demo"
        vector_store = AzureSearch(
            azure_search_endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
            azure_search_key=os.environ["AZURE_SEARCH_KEY"],
            index_name=index_name,
            embedding_function=embedding_model.embed_query,
        )
        return vector_store


